package com.small.xx.ai.spring.advisor.controller;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor;
import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.List;

import static org.springframework.ai.chat.memory.ChatMemory.CONVERSATION_ID;

/**
 * advisor绑定内存记忆案例
 * PromptChatMemoryAdvisor memory
 * 则将历史对话转换为非结构化的纯文本，直接拼接到提示词的系统文本部分，适用于对格式要求不高的简单场景
 * 能因文本混合导致上下文界限模糊，但更节省 token 消耗
 * 更适合对历史信息进行概括性引用的场景，如知识问答
 * PromptChatMemoryAdvisor 进行文本拼接
 */
@RestController
@RequestMapping("/advisor/memory/prompt")
public class MemoryPromptAdvisorController {
    private final InMemoryChatMemoryRepository chatMemoryRepository = new InMemoryChatMemoryRepository();
    private final int MAX_MESSAGES = 10;
    private final ChatClient.Builder builder;
    private ChatClient chatClient;
    private MessageWindowChatMemory messageWindowChatMemory;

    public MemoryPromptAdvisorController(ChatClient.Builder builder) {
        this.builder = builder;
        // 设置内存消息最大条数
        messageWindowChatMemory =
            MessageWindowChatMemory.builder().chatMemoryRepository(chatMemoryRepository).maxMessages(MAX_MESSAGES)
                .build();
        this.chatClient =
            builder.defaultAdvisors(PromptChatMemoryAdvisor.builder(messageWindowChatMemory).build()).build();
    }

    @GetMapping("/config")
    public String config(@RequestParam(name = "maxMessage", defaultValue = "10") int maxMessages) {
        // 设置内存消息最大条数
        messageWindowChatMemory =
            MessageWindowChatMemory.builder().chatMemoryRepository(chatMemoryRepository).maxMessages(MAX_MESSAGES)
                .build();
        this.chatClient =
            builder.defaultAdvisors(PromptChatMemoryAdvisor.builder(messageWindowChatMemory).build()).build();
        return "改为：" + maxMessages;
    }

    @GetMapping("/call")
    public String call(@RequestParam(value = "query", defaultValue = "你好，我叫xx,的外号是XX，请记住呀") String query,
        @RequestParam(value = "conversation_id", defaultValue = "XX") String conversationId) {
        return chatClient.prompt(query).advisors(a -> a.param(CONVERSATION_ID, conversationId)).call().content();
    }

    @GetMapping("/messages")
    public List<Message> messages(@RequestParam(value = "conversation_id", defaultValue = "XX") String conversationId) {
        return messageWindowChatMemory.get(conversationId);
    }
}
